import os
import pandas as pd
import json
from PyQt5.QtWidgets import (QDialog, QVBoxLayout, QLabel, QPushButton, 
                             QTableWidget, QTableWidgetItem, QMessageBox,
                             QHeaderView, QFileDialog)
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QColor, QFont
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
import numpy as np

class PredictionDialog(QDialog):
    """预测结果弹窗"""
    def __init__(self, predictions, parent=None):
        super().__init__(parent)
        self.parent = parent
        self.setWindowTitle("净空单变化预测分析")
        self.setGeometry(200, 200, 1000, 700)
        
        layout = QVBoxLayout(self)
        
        # 添加说明标签
        description = QLabel("基于当前净空单的显著变化，预测未来指数走势")
        description.setStyleSheet("font-weight: bold; color: #2c3e50;")
        layout.addWidget(description)
        
        # 创建表格显示预测结果
        self.table = QTableWidget()
        self.table.setColumnCount(8)
        self.table.setHorizontalHeaderLabels([
            "指数", "预测日期", "预测类型", "预测滞后天数", 
            "净空单变化量", "预测概率", "预测平均涨幅", "置信度"
        ])
        self.table.verticalHeader().setVisible(False)
        self.table.setEditTriggers(QTableWidget.NoEditTriggers)
        self.table.setSelectionBehavior(QTableWidget.SelectRows)
        self.table.setSelectionMode(QTableWidget.SingleSelection)
        self.table.horizontalHeader().setSectionResizeMode(QHeaderView.Interactive)
        self.table.horizontalHeader().setStretchLastSection(True)
        
        layout.addWidget(self.table)
        
        # 添加保存设置按钮
        save_settings_btn = QPushButton("保存当前设置为默认")
        save_settings_btn.setToolTip("将当前预测设置保存为默认配置")
        save_settings_btn.clicked.connect(self.save_as_default)
        layout.addWidget(save_settings_btn)
        
        # 添加导出按钮
        export_btn = QPushButton("导出预测结果为CSV")
        export_btn.clicked.connect(self.export_predictions)
        layout.addWidget(export_btn)
        
        # 添加图表按钮
        chart_btn = QPushButton("查看预测图表")
        chart_btn.clicked.connect(self.show_chart)
        layout.addWidget(chart_btn)

        # 添加评估按钮
        self.evaluate_btn = QPushButton("评估预测性能")
        self.evaluate_btn.clicked.connect(self.show_evaluation_dialog)
        layout.addWidget(self.evaluate_btn)
        
        # 存储预测结果
        self.predictions = predictions
        self.populate_table()
    
    def populate_table(self):
        """填充预测结果表格"""
        if not self.predictions:
            self.table.setRowCount(1)
        
            # 创建合并单元格
            self.table.setColumnCount(1)
            self.table.setHorizontalHeaderLabels(["预测结果详情"])
            
            # 设置行高
            self.table.setRowHeight(0, 100)
            
            # 创建文本项
            reason_item = QTableWidgetItem("无预测结果\n\n原因分析:")
            reason_item.setTextAlignment(Qt.AlignCenter | Qt.AlignVCenter)
            reason_item.setFont(QFont("Arial", 12, QFont.Bold))
            
            # 添加详细原因
            details = [
                "• 当前净空单变化量未达到设定的显著变化阈值",
                "• 可能没有足够的历史数据支持预测",
                "• 预测参数设置可能过于严格",
                "• 最近市场波动较小，无明显变化信号"
            ]
            
            reason_item.setToolTip("\n".join(details))
            self.table.setItem(0, 0, reason_item)
            return
        
                
        # 如果有预测结果，设置正常列数
        if self.table.columnCount() != 8:
            self.table.setColumnCount(8)
            self.table.setHorizontalHeaderLabels([
                "指数", "预测日期", "预测类型", "预测滞后天数", 
                "净空单变化量", "预测概率", "预测平均涨幅", "置信度"
            ])
            
        self.table.setRowCount(len(self.predictions))
        
        for row, pred in enumerate(self.predictions):
            # 获取中文名称
            chinese_name = self.parent.get_chinese_name(pred['合约'])
            
            # 指数名称
            index_item = QTableWidgetItem(chinese_name)
            index_item.setTextAlignment(Qt.AlignCenter)
            self.table.setItem(row, 0, index_item)
            
            # 预测日期
            date_item = QTableWidgetItem(str(pred['预测日期']))
            date_item.setTextAlignment(Qt.AlignCenter)
            self.table.setItem(row, 1, date_item)
            
            # 预测类型
            pred_type = "看跌" if "增大" in pred['预测类型'] else "看涨"
            type_item = QTableWidgetItem(pred_type)
            type_item.setTextAlignment(Qt.AlignCenter)
            self.table.setItem(row, 2, type_item)
            
            # 预测滞后天数
            lag_item = QTableWidgetItem(f"T+{pred['预测滞后天数']}")
            lag_item.setTextAlignment(Qt.AlignCenter)
            self.table.setItem(row, 3, lag_item)
            
            # 净空单变化量
            change_item = QTableWidgetItem(f"{pred['净空单变化量']:.2f}")
            change_item.setTextAlignment(Qt.AlignCenter)
            self.table.setItem(row, 4, change_item)
            
            # 预测概率
            prob = pred['预测概率']
            prob_item = QTableWidgetItem(f"{prob:.2%}")
            prob_item.setTextAlignment(Qt.AlignCenter)
            
            # 根据概率值设置背景色
            if prob > 0.7:
                prob_item.setBackground(QColor(255, 200, 200))  # 红色
            elif prob > 0.6:
                prob_item.setBackground(QColor(255, 235, 156))  # 黄色
            else:
                prob_item.setBackground(QColor(200, 255, 200))  # 绿色
                
            self.table.setItem(row, 5, prob_item)
            
            # 预测平均涨幅
            change_pct = pred['预测平均涨幅']
            change_pct_item = QTableWidgetItem(f"{change_pct:.2f}%")
            change_pct_item.setTextAlignment(Qt.AlignCenter)
            
            # 添加预测方法提示
            if '预测方法' in pred:
                change_pct_item.setToolTip(f"预测方法: {pred['预测方法']}")
            
            # 根据涨幅设置颜色
            if change_pct < 0:
                change_pct_item.setForeground(QColor(255, 0, 0))  # 红色
            else:
                change_pct_item.setForeground(QColor(0, 150, 0))  # 绿色
                
            self.table.setItem(row, 6, change_pct_item)
            
            # 置信度（基于样本数量）
            sample_count = pred.get('满足条件样本数', 0)
            confidence = "高" if sample_count > 50 else ("中" if sample_count > 20 else "低")
            conf_item = QTableWidgetItem(confidence)
            conf_item.setTextAlignment(Qt.AlignCenter)
            
            # 根据置信度设置背景色
            if confidence == "高":
                conf_item.setBackground(QColor(200, 255, 200))  # 绿色
            elif confidence == "中":
                conf_item.setBackground(QColor(255, 255, 200))  # 黄色
            else:
                conf_item.setBackground(QColor(255, 200, 200))  # 红色
                
            self.table.setItem(row, 7, conf_item)

            # 阈值方法
            method_item = QTableWidgetItem(pred['阈值方法'])
            method_item.setTextAlignment(Qt.AlignCenter)
            self.table.setItem(row, 8, method_item)
            
            # 样本数
            sample_count = pred.get('满足条件样本数', 0)
            sample_item = QTableWidgetItem(str(sample_count))
            sample_item.setTextAlignment(Qt.AlignCenter)
            self.table.setItem(row, 9, sample_item)

        # 调整列宽
        
        # 调整列宽
        self.table.resizeColumnsToContents()
        
        # 设置第一列宽度
        self.table.setColumnWidth(0, 100)
        
        # 设置第二列宽度
        self.table.setColumnWidth(1, 100)
        
        # 设置预测类型列宽度
        self.table.setColumnWidth(2, 80)

        # 设置阈值方法列宽度
        self.table.setColumnWidth(8, 100)
        
        # 设置样本数列宽度
        self.table.setColumnWidth(9, 80)
        
        # 设置净空单变化量列宽度
        self.table.setColumnWidth(4, 120)
        self.table.resizeColumnsToContents()
    
    def export_predictions(self):
        """导出预测结果为CSV"""
        if not self.predictions or not self.parent.output_dir:
            QMessageBox.warning(self, "警告", "没有可导出的预测结果或未选择输出目录")
            return
            
        try:
            # 创建数据框
            df = pd.DataFrame(self.predictions)
            
            # 添加中文名称列
            df['指数'] = df['合约'].apply(lambda x: self.parent.get_chinese_name(x))
            
            # 选择需要的列
            df = df[[
                '指数', '预测日期', '预测日期类型', '预测滞后天数', 
                '净空单变化量', '预测类型', '预测概率', '预测平均涨幅',
                '阈值', '阈值方法', '是否满足阈值'
            ]]
            
            # 保存文件
            output_path = os.path.join(self.parent.output_dir, '净空单预测结果.csv')
            df.to_csv(output_path, index=False, float_format='%.4f')
            
            QMessageBox.information(self, "导出成功", f"预测结果已导出到:\n{output_path}")
        except Exception as e:
            QMessageBox.critical(self, "导出失败", f"导出预测结果时出错:\n{str(e)}")
    
    def show_chart(self):
        """显示预测图表"""
        if not self.predictions:
            QMessageBox.warning(self, "无数据", "没有可显示的预测结果")
            return
            
        # 创建图表
        figure = plt.figure(figsize=(14, 10), dpi=100)  # 增加图表尺寸
        
        # 设置支持中文的字体
        plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'WenQuanYi Zen Hei']
        plt.rcParams['axes.unicode_minus'] = False
        
        # 按指数分组预测结果
        index_data = {}
        for pred in self.predictions:
            index_name = self.parent.get_chinese_name(pred['合约'])
            if index_name not in index_data:
                index_data[index_name] = {'lags': [], 'probs': [], 'changes': []}
            
            index_data[index_name]['lags'].append(pred['预测滞后天数'])
            index_data[index_name]['probs'].append(pred['预测概率'])
            index_data[index_name]['changes'].append(pred['预测平均涨幅'])
        
        # 创建图表 - 一行两列
        ax1 = figure.add_subplot(211)
        ax2 = figure.add_subplot(212, sharex=ax1)  # 共享X轴
        
        # 设置颜色
        colors = plt.cm.tab10(np.linspace(0, 1, len(index_data)))
        
        # 绘制预测概率图表
        for i, (index_name, data) in enumerate(index_data.items()):
            # 对滞后天数排序以确保正确的连接顺序
            sorted_data = sorted(zip(data['lags'], data['probs']), key=lambda x: x[0])
            lags_sorted = [d[0] for d in sorted_data]
            probs_sorted = [d[1] for d in sorted_data]
            ax1.plot(lags_sorted, probs_sorted, 'o-', color=colors[i], 
                    markersize=8, linewidth=2, label=index_name)
            
            # 添加数据标签
            for lag, prob in zip(lags_sorted, probs_sorted):
                ax1.text(lag, prob + 0.02, f"{prob:.2%}", 
                        fontsize=9, ha='center', va='bottom')
        
        ax1.set_title('未来指数走势预测概率', fontsize=14)
        ax1.set_xlabel('滞后天数', fontsize=12)
        ax1.set_ylabel('概率', fontsize=12)
        ax1.legend(loc='best')
        ax1.grid(True, linestyle='--', alpha=0.6)
        ax1.set_ylim(0, 1.1)  # 增加一点空间用于标签
        ax1.axhline(y=0.5, color='r', linestyle='--', alpha=0.5)
        
        # 绘制预测平均涨幅图表
        for i, (index_name, data) in enumerate(index_data.items()):
            # 对滞后天数排序
            sorted_data = sorted(zip(data['lags'], data['changes']), key=lambda x: x[0])
            lags_sorted = [d[0] for d in sorted_data]
            changes_sorted = [d[1] for d in sorted_data]
            
            # 使用条形图显示
            ax2.bar([d + i*0.2 for d in lags_sorted], changes_sorted, 
                    width=0.2, color=colors[i], label=index_name)
            
            # 添加数据标签
            for lag, change in zip(lags_sorted, changes_sorted):
                color = 'red' if change < 0 else 'green'
                ax2.text(lag + i*0.2, change + (0.1 if change >=0 else -0.5), 
                        f"{change:.2f}%", fontsize=9, ha='center', va='bottom', color=color)
        
        ax2.set_title('预测平均涨幅', fontsize=14)
        ax2.set_xlabel('滞后天数', fontsize=12)
        ax2.set_ylabel('涨幅百分比(%)', fontsize=12)
        ax2.legend(loc='best')
        ax2.grid(True, linestyle='--', alpha=0.6, axis='y')
        ax2.axhline(y=0, color='gray', linestyle='-', alpha=0.5)
        
        # 设置x轴刻度
        all_lags = sorted(set(lag for data in index_data.values() for lag in data['lags']))
        ax1.set_xticks(all_lags)  # 设置x轴刻度位置
        ax1.set_xticklabels(all_lags)  # 设置x轴标签
        ax2.set_xticks(all_lags)  # 设置x轴刻度位置
        ax2.set_xticklabels(all_lags)
        
        # 调整布局
        figure.tight_layout(pad=3.0)  # 增加内边距
        figure.subplots_adjust(top=0.9, hspace=0.3)  # 增加子图间距
        
        # 创建自定义图表窗口
        win = QDialog(self)
        win.setWindowTitle("预测图表")
        win.setGeometry(100, 100, 1200, 800)
        layout = QVBoxLayout(win)
        
        # 将图表嵌入到窗口中
        canvas = FigureCanvas(figure)
        layout.addWidget(canvas)
        
        # 添加工具栏
        toolbar = NavigationToolbar(canvas, win)
        layout.addWidget(toolbar)
        
        # 添加导出按钮
        export_btn = QPushButton("导出图表")
        def export_handler():
            file_path, _ = QFileDialog.getSaveFileName(
                win, "保存图表", "", 
                "PNG图片 (*.png);;JPEG图片 (*.jpg);;PDF文件 (*.pdf);;SVG文件 (*.svg)"
            )
            if not file_path:
                return
                
            try:
                if file_path.endswith('.png'):
                    figure.savefig(file_path, dpi=300, format='png')
                elif file_path.endswith('.jpg') or file_path.endswith('.jpeg'):
                    figure.savefig(file_path, dpi=300, format='jpeg')
                elif file_path.endswith('.pdf'):
                    figure.savefig(file_path, format='pdf')
                elif file_path.endswith('.svg'):
                    figure.savefig(file_path, format='svg')
                else:
                    file_path += '.png'
                    figure.savefig(file_path, dpi=300, format='png')
                    
                QMessageBox.information(win, "导出成功", f"图表已保存到:\n{file_path}")
            except Exception as e:
                QMessageBox.critical(win, "导出失败", f"保存图表时出错:\n{str(e)}")
        export_btn.clicked.connect(export_handler)
        layout.addWidget(export_btn)
        
        win.exec_()

    def show_evaluation_dialog(self):
        dialog = self.parent.PredictionEvaluationDialog(self.parent)
        dialog.exec_()
    
    def save_as_default(self):
        """将当前预测设置保存为默认设置"""
        if not self.predictions:
            QMessageBox.warning(self, "无设置", "没有可保存的预测设置")
            return
            
        try:
            # 确保父窗口有预测设置
            if hasattr(self.parent, 'prediction_settings') and self.parent.prediction_settings:
                # 使用QSettings保存到配置文件
                settings = QSettings("QuantAnalysis", "IncrementalDataProcessor")
                settings.setValue("prediction_settings", json.dumps(self.parent.prediction_settings))
                
                QMessageBox.information(self, "保存成功", "当前预测设置已保存为默认设置！")
            else:
                QMessageBox.warning(self, "无设置", "没有找到可保存的预测设置")
        except Exception as e:
            QMessageBox.critical(self, "保存失败", f"保存设置时出错:\n{str(e)}")